import json
import os
import time

import requests

from .leetcode_operation import LeetCodeOperation


class MemoryDistribution(LeetCodeOperation):
    def __init__(self, requests_manager, base_url="https://leetcode.cn/submissions/api/memory_distribution/", **params):
        super().__init__(requests_manager, base_url)
        self.endpoint = 'memory_distribution'

    def parse_graphql_data(self, data, submission_id):
        # Extract the necessary fields from the GraphQL response
        if not data:
            return None
        return {
            'merged_distribution': data.get('merged_distribution', []),
            'lang_display': data.get('lang_display', ''),
            'observerId': submission_id
        }

    def fetch_all_data(self, request_interval=1, folder_path='MemoryDistribution'):
        submission_list_dir = 'SubmissionList'
        os.makedirs(os.path.join(self.directory, folder_path), exist_ok=True)

        for slug_file in os.listdir(os.path.join(self.directory, submission_list_dir)):
            save_path = os.path.join(folder_path, f"{slug_file}")
            if os.path.exists(os.path.join(self.directory, save_path)):
                continue
            
            self.logger.info(f"slug_file: {slug_file}")
            slug_data = self.load_json(os.path.join(submission_list_dir, slug_file))
            submission_id = slug_data['submissions'][0]['id'] if slug_data['submissions'] else None
            
            if submission_id:
                memory_data = self.fetch_data(submission_id, method='GET')
                parsed_data = self.parse_graphql_data(memory_data, submission_id)
                if parsed_data:
                    self.save_json(parsed_data, save_path)

                time.sleep(request_interval)
            else:
                os.unlink(os.path.join(submission_list_dir, slug_file))

    def execute(self):
        data = self.fetch_data(self.endpoint, method='POST')
        parsed_data = self.parse_graphql_data(data)
        self.save_json(parsed_data, self.filename)

